package com.hiss.basic.starter.core.utils.tree;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.collection.ListUtil;
import cn.hutool.core.util.ObjectUtil;
import lombok.AccessLevel;
import lombok.NoArgsConstructor;

import java.util.*;
import java.util.function.Predicate;
import java.util.stream.Collectors;

/**
 * 树形构建工具
 */
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public class TreeBuilder {

    public static <T, W> List<TreeNode<T, W>> build(List<TreeNode<T, W>> nodeList) {
        return build(nodeList, parentId -> ObjectUtil.equal(0L, parentId));
    }

    public static <T, W> List<TreeNode<T, W>> build(List<TreeNode<T, W>> nodeList, Predicate<T> rootPredicate) {
        return build(nodeList, rootPredicate, null);
    }

    public static <T, W> List<TreeNode<T, W>> build(List<TreeNode<T, W>> nodeList, Predicate<T> rootPredicate, Comparator<TreeNode<T, W>> comparator) {
        if (CollectionUtil.isEmpty(nodeList)) {
            return ListUtil.toList();
        }

        Map<T, TreeNode<T, W>> nodeMap = nodeList.stream().collect(Collectors.toMap(TreeNode::getId, node -> node));
        List<TreeNode<T, W>> rootNodeList = ListUtil.toList();

        for (TreeNode<T, W> node : nodeList) {
            T parentId = node.getParentId();
            if (rootPredicate.test(parentId)) {
                rootNodeList.add(node);
            } else {
                TreeNode<T, W> parentNode = nodeMap.get(parentId);
                if (ObjectUtil.isNotNull(parentNode)) {
                    if (CollectionUtil.isEmpty(parentNode.getChildren())) {
                        parentNode.setChildren(new ArrayList<>());
                    }
                    parentNode.getChildren().add(node);
                } else {
                    rootNodeList.add(node);
                }
            }
        }

        if (CollectionUtil.isNotEmpty(rootNodeList) && ObjectUtil.isNotNull(comparator)) {
            rootNodeList.sort(comparator);
            for (TreeNode<T, W> node : rootNodeList) {
                if (ObjectUtil.isNotNull(node.getChildren())) {
                    node.getChildren().sort(comparator);
                }
            }
        }
        return rootNodeList;
    }

    /**
     * 获取最后一级叶子节点id列表
     *
     * @param nodeList
     * @param <T>
     * @param <W>
     * @return
     */
    public static <T, W> List<T> findLastLeafNodeIds(List<TreeNode<T, W>> nodeList) {
        List<T> leafNodeIds = new ArrayList<>();
        if (CollectionUtil.isEmpty(nodeList)) {
            return leafNodeIds;
        }

        Stack<TreeNode<T, W>> stack = new Stack<>();
        stack.addAll(nodeList);

        while (!stack.isEmpty()) {
            TreeNode<T, W> node = stack.pop();
            if (CollectionUtil.isEmpty(node.getChildren())) {
                // 如果是叶子节点，添加到结果列表
                leafNodeIds.add(node.getId());
            } else {
                // 如果不是叶子节点，将子节点压入栈中
                stack.addAll(node.getChildren());
            }
        }
        return leafNodeIds;
    }
}
